import os


def write_list(ds_path):
    """
    write list.txt in the ds_path dataset. list.txt is especially for GDN model
    :param ds_path: dataset path
    :return: None
    """
    if not os.path.exists(f'{ds_path}/list.txt'):
        train_csv_path = os.listdir(f'{ds_path}/train/')
        with open(f'{ds_path}/train/{train_csv_path[0]}', 'r') as train_csv_file:
            # this is what we're going to put in the list.txt: all the column info in the train dataset
            column_list = train_csv_file.readline().strip().split(',')
        with open(f'{ds_path}/list.txt', 'w') as list_txt:
            for column in column_list:
                list_txt.write(f'{column}\n')


def gdn_pre(data_dir):
    datasets = os.listdir(data_dir)
    for dataset in datasets:
        if dataset in ['SWAT', 'WADI']:
            # SWAT and WADI dataset don't have sub-dataset
            if os.path.isdir(os.path.join(data_dir, dataset)):
                write_list(os.path.join(data_dir, dataset))
        elif dataset in ['MSL', 'SMAP', 'SMD']:
            # these three datasets have sub-datasets
            sub_datasets = os.listdir(os.path.join(data_dir, dataset))
            for sub_dataset in sub_datasets:
                # Because there are some strange files in the dir, we need to judge whether it's a directory or not
                if os.path.isdir(os.path.join(data_dir, dataset, sub_dataset)):
                    write_list(os.path.join(data_dir, dataset, sub_dataset))
